import CoolProp.CoolProp as CP
import matplotlib.pyplot as plt

print('压焓图及循环绘制 版本1.0\n')
ref = input('输入冷媒名称(R290、R32、R410A、R404A、R134a、CO2等): ')

#画饱和液态、气态线  -40℃对应压力~临界压力
P_crit = CP.PropsSI(ref, 'Pcrit') #获取临界压力 Pa
P_max = P_crit+2e6 #绘图上限 Pa
if ref == 'CO2':
    P_max = P_crit+10e6 #对于CO2另做调整
P_min = CP.PropsSI('P', 'T', -35+273.15, 'Q', 0, ref) #-35℃对应压力 Pa

interval = 100 #划分区间个数
h_liquid = [1]*(interval+1)
h_vapor = [1]*(interval+1)
P = [1]*(interval+1)
dP = (P_crit-P_min)/interval
for i in range(interval+1):
    h_liquid[i] = CP.PropsSI('H', 'P', P_min+dP*i, 'Q', 0, ref)/1e3 #获取液态饱和线上焓值 kJ/kg
    h_vapor[i] = CP.PropsSI('H', 'P', P_min+dP*i, 'Q', 1, ref)/1e3 #获取气态饱和线上焓值 kJ/kg
    P[i] = (P_min+dP*i)/1e6
plt.figure(figsize=(10, 6.18))
plt.plot(h_liquid, P, color='black', linewidth=1.5) #液态线
plt.plot(h_vapor, P, color='black', linewidth=1.5) #气态线
#plt.xlim(min(h_liquid)-30, max(h_vapor)+200)
plt.ylim(P_min/1e6, P_max/1e6)
plt.yscale('log')
plt.xlabel('Enthalpy/kJ·kg-1')
plt.ylabel('Pressure/MPa')
plt.title(ref+' p-h diagram',y = 1.03) #增加标题与图表的距离1.03

#画等温线
interval = 100 #划分区间个数
T = -30+273.15 #最低温度线 K
T_crit = CP.PropsSI(ref, 'Tcrit') #临界温度 K
dT = 10 #温度间隔
while T < T_crit+40: #最高温度线为临界温度+40℃
    h = [1]*(interval+1)
    P = [1]*(interval+1)
    dP = (P_max-P_min)/interval
    for i in range(interval+1):
        try: #R410A、R404A混合物在靠近临界点两相区计算焓值会出错，因此这里在报错时用上一个点替代，整体上不影响准确性
            h[i] = CP.PropsSI('H', 'T', T, 'P', P_min+dP*i, ref)/1e3
            P[i] = (P_min+dP*i)/1e6
        except:
            h[i] = h[i-1]
            P[i] = P[i-1]
    if T < T_crit: #插入两个饱和点
        h_liquid1 = CP.PropsSI('H', 'T', T, 'Q', 0, ref)/1e3
        P_liquid1 = CP.PropsSI('P', 'T', T, 'Q', 0, ref)/1e6
        h_vapor1 = CP.PropsSI('H', 'T', T, 'Q', 1, ref)/1e3
        P_vapor1 = CP.PropsSI('P', 'T', T, 'Q', 1, ref)/1e6
        for i in range(interval+1):
            if h[i] < h_vapor1:
                P.insert(i, P_liquid1)
                h.insert(i, h_liquid1)
                P.insert(i, P_vapor1)
                h.insert(i, h_vapor1)
                plt.text(h[i+1]+10, P[i+1], str(round(T-273.15))+'℃') #取靠近饱和液体线数据加标签
                break
    plt.plot(h, P, color='red', linewidth=0.5) #红色线、线宽0.5
    T = T+dT

#画等干度线
interval = 100 #划分区间个数
for Q in [0.2, 0.4, 0.6, 0.8]: #指定干度
    h = [1]*(interval+1)
    P = [1]*(interval+1)
    dP = (P_crit-P_min)/interval
    for i in range(interval+1):
        h[i] = CP.PropsSI('H', 'P', P_min+dP*i, 'Q', Q, ref)/1e3
        P[i] = (P_min+dP*i)/1e6
    plt.text(h[1], P[1], Q) #取第二个数加标签
    plt.plot(h, P, color='black', linewidth=0.5)

#画等熵线
ds = 0.2
s_min = CP.PropsSI('S', 'P', P_min, 'Q', 0, ref)/1e3
s_min = round(s_min, 1) #round为四舍五入，这一步是为了取整1位小数
s_max = CP.PropsSI('S', 'P', P_min, 'Q', 1, ref)/1e3
s_max = round(s_max, 1)+ds
if (s_max-s_min)/6 < ds: #0.2跨度太大情况下改为0.1
    ds = 0.1
s = s_min
interval = 100
while s <= s_max:
    h = [1]*(interval+1)
    P = [1]*(interval+1)
    dP = (P_max-P_min)/interval
    for i in range(interval+1):
        try: #CO2在超临界的地方计算失败，采取方式控制
            h[i] = CP.PropsSI('H', 'P', P_min+dP*i, 'S', s*1e3, ref)/1e3
            P[i] = (P_min+dP*i)/1e6
        except:
            h[i] = h[i-1]
            P[i] = P[i-1]
    plt.text(h[-15], P[-15], round(s,2))
    plt.plot(h, P, color='blue', linewidth=0.5)
    s = s+ds
h_es_max = max(h) #用于限定等密度曲线的范围，控制最大焓

#画等密度线
drho = 5
rho_min = CP.PropsSI('D', 'P', P_min, 'Q', 1, ref)
rho_min = round(rho_min) #round为四舍五入，这一步是为了取整1位小数
rho_max = CP.PropsSI('D', 'P', P_min, 'Q', 0, ref)
rho_max = round(rho_max)-drho
rho = rho_min
interval = 100
j = 1
while rho <= rho_max:
    h = [1]*(interval+1)
    P = [1]*(interval+1)
    dP = (P_max-P_min)/interval
    for i in range(interval+1):
        h[i] = CP.PropsSI('H', 'P', P_min+dP*i, 'D', rho, ref)/1e3
        if h[i] > h_es_max:
            h = h[:i-1]
            P = P[:i-1]
            break
        P[i] = (P_min+dP*i)/1e6
    if len(h) > 0: #防止h为空，计算R1234yf时发现
        plt.text(h[-1], P[-1], rho)
        plt.plot(h, P, color='green', linewidth=0.5)
    rho = rho+drho*j*j #用x^2曲线
    j = j+1

#循环绘制，两种模式，一种是4点简易循环，一种是根据输入的压力、温度数组绘制循环，不限点数
index = input("循环绘制 不绘制直接回车 简易循环(4点)输入1 复杂循环(多点)输入2 ")
try:
    index = int(index)
    if index == 1:
        pt = list(map(float, input("请输入以空格分隔的吸气压力/MPa、排气压力/MPa、吸气温度/℃、阀前温度/℃、排气温度/℃: ").split()))
        P_cy = [pt[0], pt[1], pt[1], pt[0], pt[0]]
        h_cy = [1]*5
        h_cy[0] = CP.PropsSI('H', 'P', pt[0]*1e6, 'T', pt[2]+273.15, ref)/1e3 #吸气比焓
        h_cy[1] = CP.PropsSI('H', 'P', pt[1]*1e6, 'T', pt[4]+273.15, ref)/1e3 #排气比焓
        h_cy[2] = CP.PropsSI('H', 'P', pt[1]*1e6, 'T', pt[3]+273.15, ref)/1e3 #排气比焓
        h_cy[3] = h_cy[2]
        h_cy[4] = h_cy[0]

    elif index == 2:
        P_cy = list(map(float, input("请输入以空格分隔的压力/MPa(绘图时将按顺序连接,自动闭环): ").split()))
        P_cy.append(P_cy[0])
        P_cy_c = [x*1e6 for x in P_cy]
        t_cy = list(map(float, input("请输入以空格分隔的温度/℃(绘图时将按顺序连接,自动闭环): ").split()))
        t_cy.append(t_cy[0])
        t_cy_c = [x+273.15 for x in t_cy]
        h_cy = CP.PropsSI('H', 'P', P_cy_c, 'T', t_cy_c, ref)/1e3
    plt.plot(h_cy, P_cy, color='magenta', linewidth=1.5)
except:
    print('跳过循环绘制')

plt.show()
input('check')